import asyncio
import json
import ssl
from urllib.parse import urlencode
from datetime import datetime

# 引入 mcp 客户端库
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

# --- 全局禁用SSL证书验证 ---
try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    pass
else:
    ssl._create_default_https_context = _create_unverified_https_context
# --- ---


async def find_available_tickets_for_train(api_key, train_codes, departure_cities, arrival_cities, query_date=None, *, group_name=None):
    # 为一个“查询组”并行查询余票情况并返回有余票的结果条目。
    # 行为概述
    # - 在单独的 MCP 连接（通过 streamablehttp_client）中为该查询组创建并复用一个 ClientSession。
    # - 对每一对出发城市/到达城市（departure_cities × arrival_cities，排除相同城市对）并发发起查询，使用 semaphore 限制并发数（CONCURRENCY_LIMIT = 20）。
    # - 使用 session.call_tool("search", {"date": query_date, "fromCity": dep, "toCity": arr}) 获取原始查询结果，解析返回的 JSON，筛选出与所需车次匹配且存在可用席位的记录。
    # - 返回每个发现有余票的条目（字典），包括查询组名、日期、车次、出发/到达地、出发/到达时间以及具体可用席位与数量（以字符串形式保留原始数值）。
    # 参数
    # - api_key (str): 用于 MCP 连接的 API Key，会被拼接到 base_url 的查询参数中。
    # - train_codes (str | Iterable[str] | None): 要关注的车次，可以是单个车次字符串或可迭代的车次集合（list/tuple/set 等）。函数会对元素 str(...).strip() 并去重。若为空或 None，则等同于不关注任何车次（通常会导致无匹配）。
    # - departure_cities (Iterable[str] | None): 出发城市列表/可迭代对象。支持 None（视为空列表）。
    # - arrival_cities (Iterable[str] | None): 到达城市列表/可迭代对象。支持 None（视为空列表）。
    # - query_date (str | None): 查询日期，格式应为 'YYYY-MM-DD'，若为 None 则使用当前日期 datetime.now().strftime('%Y-%m-%d')。
    # - group_name (str | None, keyword-only): 可选的“查询组名”，仅用于日志前缀和在结果中标记来源组，便于多组并行查询时区分。
    # 返回值
    # - List[dict]: 列表中每个元素表示一次匹配到有余票的查询结果，字典示例结构如下：
    #     {
    #         "查询组": <group_name or None>,
    #         "日期": <query_date str>,
    #         "车次": <车次字符串>,
    #         "出发地": <出发城市字符串>,
    #         "到达地": <到达城市字符串>,
    #         "出发时间": <出发时间字符串或空字符串>,
    #         "到达时间": <到达时间字符串或空字符串>,
    #         "余票席位": {
    #                 "商务座": "3",           # 原始数值以字符串形式保留
    #                 "一等座": "有",          # 只要不在无效集合即被视为可用并保留原样（已 strip）
    #                 ...
    # - 若没有任何有效城市组合、没有匹配到车次或未发现有余票条目，则返回空列表 []。
    # - 如果在建立连接或初始化 MCP 会话时发生异常，函数会打印日志并返回 []。
    # 实现细节与约定
    # - 并发限制为 20（CONCURRENCY_LIMIT = 20），通过 asyncio.Semaphore 控制。
    # - seat_map 定义了中文席位名称到票数据字段名（如 '商务座' -> 'swz_num'）的映射，用于从票列表中抽取票数。
    # - 对“无效票数”的判定集合为 {"无", "--", "0", "", "-"}，这些值会被视为无票并排除；其他非 None 值会作为可用票保留（以 str 形式）。
    # - 车次匹配使用 strip() 后的精确匹配（不做模糊或正则匹配）。
    # - 从工具返回值解析时，函数期望 result.content 为列表且其首项 .text 为 JSON 字符串（即 json.loads(result.content[0].text) 可得到票信息列表），并从每条记录中读取键如 '车次'、'票列表'、'出发时间' / '始发时间' / '发车时间' / '出发时刻'、'到达时间' / '终到时间' / '到站时间' / '到达时刻' 等。
    # - 函数会打印若干日志（使用 print），包括任务开始/完成、每次发现余票或异常的记录，查询组名（若提供）会作为日志前缀。
    # - 对单个城市对的查询若发生异常，会捕获并记录但不会抛出，其他并发任务会继续执行；最终会汇总所有成功发现的结果返回。
    # 示例（伪代码）
    # - 调用示例：
    #         await find_available_tickets_for_train(
    #                 api_key="MYKEY",
    #                 train_codes=["G101", "D202"],
    #                 departure_cities=["北京", "天津"],
    #                 arrival_cities=["上海", "杭州"],
    #                 query_date="2025-09-20",
    #                 group_name="早班次关注"
    # - 返回示例：
    #         [
    #             {
    #                 "查询组": "早班次关注",
    #                 "日期": "2025-09-20",
    #                 "车次": "G101",
    #                 "出发地": "北京",
    #                 "到达地": "上海",
    #                 "出发时间": "08:00",
    #                 "到达时间": "12:30",
    #                 "余票席位": {"二等座": "5", "无座": "有"}
    #             },
    #             ...
    #         ]
    # 注意事项
    # - 函数依赖外部组件 streamablehttp_client 与 ClientSession 的具体实现以及 session.call_tool 的行为，若这些组件的接口或返回格式不同，需相应调整解析逻辑。
    # - 日志使用 print，若在生产环境需要更细粒度控制或持久化日志，请替换为合适的 logging 机制。
    """
    为一个“查询组”并行查询余票情况，并受并发数限制。
    增加了 group_name 参数用于日志区分。
    """
    # --- 并发控制器 ---
    CONCURRENCY_LIMIT = 20
    semaphore = asyncio.Semaphore(CONCURRENCY_LIMIT)
    
    # --- 日志前缀 ---
    log_prefix = f"[{group_name}] " if group_name else ""

    seat_map = {
        '商务座': 'swz_num', '一等座': 'zy_num', '二等座': 'ze_num', '高级软卧': 'gr_num',
        '软卧': 'rw_num', '动卧': 'dw_num', '硬卧': 'yw_num', '软座': 'rz_num',
        '硬座': 'yz_num', '无座': 'wz_num'
    }

    if isinstance(train_codes, (list, tuple, set)):
        wanted_trains = {str(t).strip() for t in train_codes if t}
    else:
        wanted_trains = {str(train_codes).strip()} if train_codes else set()

    if not query_date:
        query_date = datetime.now().strftime('%Y-%m-%d')
    
    print(f"{log_prefix}任务开始：查询车次 【{', '.join(sorted(wanted_trains))}】 在日期 【{query_date}】 的余票情况。")

    # base_url = "https://server.smithery.ai/@other-blowsnow/mcp-server-chinarailway/mcp"
    base_url = "http://localhost:3000/sse"
    params = {"api_key": api_key}
    url_with_auth = f"{base_url}?{urlencode(params)}"
    
    async def query_single_pair(session, dep_city, arr_city):
        """在同一个 session 中查询单个城市对，并受 semaphore 控制。"""
        async with semaphore:
            # print(f"{log_prefix}  [任务启动] 查询: {dep_city} -> {arr_city}") # 这条日志太频繁，可以注释掉
            try:
                result = await session.call_tool(
                    "search",
                    {"date": query_date, "fromCity": dep_city, "toCity": arr_city}
                )
                
                ticket_data = []
                if isinstance(result.content, list) and len(result.content) > 0:
                    ticket_data = json.loads(result.content[0].text)
                
                for train_info in ticket_data:
                    actual_train_no = train_info.get("车次")
                    if actual_train_no and actual_train_no.strip() in wanted_trains:
                        ticket_list = train_info.get("票列表", {})
                        
                        invalid_vals = {"无", "--", "0", "", "-"}
                        available_seats = {
                            seat_name: str(raw).strip()
                            for seat_name in seat_map.keys()
                            if (raw := ticket_list.get(seat_name, {}).get("num")) is not None
                            and str(raw).strip() not in invalid_vals
                        }
                        
                        if available_seats:
                            start_time, end_time = None, None
                            start_keys = ['出发时间', '始发时间', '发车时间', '出发时刻']
                            end_keys = ['到达时间', '终到时间', '到站时间', '到达时刻']
                            for k in start_keys:
                                if v := train_info.get(k): start_time = str(v).strip(); break
                            for k in end_keys:
                                if v := train_info.get(k): end_time = str(v).strip(); break

                            ticket_info = {
                                "查询组": group_name, # 在结果中也加入组名
                                "日期": query_date, "车次": actual_train_no, "出发地": dep_city,
                                "到达地": arr_city, "出发时间": start_time or "", "到达时间": end_time or "",
                                "余票席位": available_seats
                            }
                            print(f"{log_prefix}  --> ✅ [查询成功] {dep_city} -> {arr_city} | 车次 {actual_train_no} 发现余票")
                            return ticket_info
                        else:
                            # print(f"{log_prefix}  --> 🟡 [已找到但无票] {dep_city} -> {arr_city} | 车次 {actual_train_no} 当前无余票。")
                            return None
                return None
            except Exception as e:
                print(f"{log_prefix}  --> 🚨 [查询异常] {dep_city} -> {arr_city}: {e}")
            return None

    all_results = []
    try:
        # 注意：每个查询组都会建立一个独立的 MCP 连接，这更稳定
        async with streamablehttp_client(url_with_auth) as (read, write, _):
            async with ClientSession(read, write) as session:
                await session.initialize()
                
                tasks = []
                # 确保 departure_cities 和 arrival_cities 是列表，避免 None 导致错误
                departure_cities = departure_cities or []
                arrival_cities = arrival_cities or []
                for dep_city in departure_cities:
                    for arr_city in arrival_cities:
                        if dep_city == arr_city: continue
                        tasks.append(query_single_pair(session, dep_city, arr_city))
                
                if not tasks:
                    print(f"{log_prefix}没有任何有效的城市组合需要查询。")
                    return []
                    
                # print(f"{log_prefix}准备并发执行 {len(tasks)} 个城市对查询任务...")
                results_from_gather = await asyncio.gather(*tasks, return_exceptions=True)
                
                for res in results_from_gather:
                    if isinstance(res, Exception):
                        print(f"{log_prefix}  --> 🚨 [任务执行异常]: {res}")
                    elif res:
                        all_results.append(res)
        print(f"{log_prefix}任务完成。")
    except Exception as e:
        print(f"{log_prefix}❌ 连接或初始化 MCP 会话失败: {type(e).__name__} - {e}")
        return []

    return all_results

# --- 使用示例 ---
async def main():
    my_api_key = "d0ff163b-d48f-4f21-8335-e78cff4342c5"
    target_date = "2025-10-01"

    # ==================== 在这里定义您所有的查询组 ====================
    query_groups = [
        {
            # allStations["潮汕", "普宁", "陆丰", "惠东南", "深圳北", "东莞南", "博罗北", "河源北", "龙川西", "和平北", "龙南东", "赣州西", "新干东", "南昌西", "庐山", "黄梅东", "太湖南", "安庆", "铜陵", "无为", "合肥南"]
            "name": "深圳北",
            "trains": ["G2691", "G2694"],
            "departures": [ "惠东南", "深圳北"],
            "arrivals": ["太湖南", "安庆"]
        },
    #   {
    #         # allStations["深圳", "河源东", "信丰西", "赣州西", "吉安西", "南昌西", "德安", "庐山", "黄梅东", "太湖南", "安庆", "池州", "铜陵", "芜湖", "马鞍山东", "南京南", "镇江", "常州", "苏州", "阳澄湖", "上海虹桥"]
    #         "name": "深圳-太湖南1",
    #         "trains": ["G2790"],
    #         "departures": [ "深圳", "河源东", "信丰西"],
    #         "arrivals": ["德安", "庐山", "黄梅东", "太湖南", "安庆", "池州", "铜陵"]
    #     },
    #   {
    #         # allStations["深圳北", "河源东", "龙川西", "龙南东", "赣州西", "吉安西", "丰城东", "南昌西", "共青城", "庐山", "宿松东", "太湖南", "安庆西", "桐城东", "肥西", "合肥南"]
    #         "name": "深圳-太湖南2",
    #         "trains": ["G636"],
    #         "departures": ["深圳北", "河源东", "龙川西"],
    #         "arrivals": [ "庐山", "宿松东", "太湖南", "安庆西", "桐城东"]
    #     },
    #   {
    #         # allStations["深圳", "仲恺", "惠州北", "定南南", "赣州西", "吉安西", "南昌西", "共青城", "庐山", "宿松东", "太湖南", "安庆", "池州", "铜陵", "芜湖", "南京南"]

    #         "name": "深圳-太湖南3",
    #         "trains": ["G2782"],
    #         "departures": ["深圳", "仲恺", "惠州北"],
    #         "arrivals": [  "共青城", "庐山", "宿松东", "太湖南", "安庆", "池州", "铜陵", "芜湖", "南京南"]
    #     },
    #   {
    #         # ["深圳", "惠州北", "龙川西", "赣州西", "吉安西", "峡江", "樟树东", "丰城东", "南昌西", "永修", "德安", "庐山", "太湖南", "潜山", "合肥南"]

    #         "name": "深圳-太湖南4",
    #         "trains": ["G2796"],
    #         "departures": ["深圳", "惠州北", "龙川西"],
    #         "arrivals": ["永修", "德安", "庐山", "太湖南", "潜山", "合肥南"]
    #     },
    #   {
    #         # allStations["深圳北", "东莞南", "惠州北", "河源东", "信丰西", "赣州西", "新干东", "南昌西", "德安", "庐山", "黄梅东", "宿松东", "潜山", "安庆", "池州", "铜陵北", "巢湖东", "合肥南", "淮南东", "蚌埠南", "宿州东"]
    #         "name": "深圳-潜山",
    #         "trains": ["G2698"],
    #         "departures": ["深圳北", "东莞南", "惠州北", "河源东"],
    #         "arrivals": [ "新干东", "南昌西", "德安", "庐山", "黄梅东", "宿松东", "潜山", "安庆", "池州", "铜陵北",]
    #     },

        # {
        #     #allStations ["广州东", "新塘", "增城", "惠州北", "河源东", "赣州西", "吉安西", "吉水西", "南昌西", "永修", "共青城", "庐山", "太湖南", "合肥南"]
        #     "name": "广州东",
        #     "trains": ["G3082"],
        #     "departures": [ "广州东", "新塘", "增城", "惠州北", "河源东"],
        #     "arrivals": [   "永修", "共青城", "庐山", "太湖南", "合肥南"]
        # },

        # {
        #     # allStations["北京丰台", "任丘", "深州", "衡水", "聊城", "菏泽", "曹县", "亳州", "阜阳", "合肥", "庐江", "桐城", "怀宁", "天柱山", "太湖", "宿松"]
        #     "name": "北京",
        #     "trains": ["K1071"],
        #     "departures": [ "北京丰台", "任丘", "深州"],
        #     "arrivals": [  "天柱山", "太湖", "宿松"]
        # },
    ]
    # =================================================================

    if not my_api_key or "YOUR_API_KEY" in my_api_key:
        print("❌ 错误：请在 main 函数中设置您的 my_api_key。")
        return

    # --- 并行执行所有查询组 ---
    tasks = []
    for i, group in enumerate(query_groups):
        # 创建一个协程任务，并传递所需的所有参数
        task = find_available_tickets_for_train(
            api_key=my_api_key,
            train_codes=group.get("trains"),
            departure_cities=group.get("departures"),
            arrival_cities=group.get("arrivals"),
            query_date=target_date,
            group_name=group.get("name", f"查询组-{i+1}") # 如果没有name字段，就给一个默认名
        )
        tasks.append(task)
    
    print(f"检测到 {len(tasks)} 个查询组，准备并行启动所有任务...\n")
    
    # 使用 asyncio.gather 并行运行所有查询组的任务
    # results_of_all_groups 会是一个列表的列表，例如： [[group1_res1, group1_res2], [group2_res1], []]
    results_of_all_groups = await asyncio.gather(*tasks)

    # --- 汇总并打印结果 ---
    final_results = []
    for group_result_list in results_of_all_groups:
        final_results.extend(group_result_list)

    print("\n\n---------- ✅ 所有查询任务完成，结果汇总如下 ----------")
    if final_results:
        # 按查询组和出发地排序，使结果更清晰
        sorted_results = sorted(final_results, key=lambda x: (x.get('查询组', ''), x['出发地'], x['车次']))
        print(json.dumps(sorted_results, ensure_ascii=False, indent=4))
    else:
        print("所有查询组合均未发现符合条件的余票。")

if __name__ == '__main__':
    asyncio.run(main())